import torch.nn as nn
import torch
from collections import OrderedDict


class Convention(nn.Module):
    def __init__(self,in_channels,out_channels,conv_size,conv_stride,padding,need_bn = True):

        """
        这边对Conv2d进行一个封装，参数一致
        但是多加了LeakReLU,和归一化，原因不多说了
        :param in_channels:
        :param out_channels:
        :param conv_size:
        :param conv_stride:
        :param padding:
        :param need_bn:
        """

        super(Convention,self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, conv_size, conv_stride, padding, bias=False if need_bn else True)
        self.leaky_relu = nn.LeakyReLU()
        self.need_bn = need_bn
        if need_bn:
            self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        return self.bn(self.leaky_relu(self.conv(x))) if self.need_bn else self.leaky_relu(self.conv(x))

    def weight_init(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                torch.nn.init.kaiming_normal_(m.weight.data)
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()


class BackboneNet(nn.Module):
    """
    骨干网络，因为那个论文中也提到了预训练的概念
    那么这个预训练其实是说训练这个骨干网络，而这个
    网络的话其实是7x7x30的前半部分
    那个yolo是24卷积+2个全连接得到7x7x1024之后flatten4096
    最后变成7x7x30，然后就是NMS，预训练需要先训练一个
    分类的网络，所以这部分是不一样的
    """
    def __init__(self):
        super(BackboneNet,self).__init__()

        """
        用于特征提取的16个卷积
        """
        self.Conv_Feature = nn.Sequential(
            Convention(3, 64, 7, 2, 3),
            nn.MaxPool2d(2, 2),
            Convention(64, 192, 3, 1, 1),
            nn.MaxPool2d(2, 2),
            Convention(192, 128, 1, 1, 0),
            Convention(128, 256, 3, 1, 1),
            Convention(256, 256, 1, 1, 0),
            Convention(256, 512, 3, 1, 1),
            nn.MaxPool2d(2, 2),
            Convention(512, 256, 1, 1, 0),
            Convention(256, 512, 3, 1, 1),
            Convention(512, 256, 1, 1, 0),
            Convention(256, 512, 3, 1, 1),
            Convention(512, 256, 1, 1, 0),
            Convention(256, 512, 3, 1, 1),
            Convention(512, 256, 1, 1, 0),
            Convention(256, 512, 3, 1, 1),
            Convention(512, 512, 1, 1, 0),
            Convention(512, 1024, 3, 1, 1),
            nn.MaxPool2d(2, 2),
        )
        self.Conv_Semanteme = nn.Sequential(
            Convention(1024, 512, 1, 1, 0),
            Convention(512, 1024, 3, 1, 1),
            Convention(1024, 512, 1, 1, 0),
            Convention(512, 1024, 3, 1, 1),
        )
